package com.xyzwps.collections.set;

import com.xyzwps.collections.Range;
import com.xyzwps.collections.Utils;

import java.util.Arrays;
import java.util.Comparator;
import java.util.Iterator;
import java.util.Objects;

/**
 * 基于跳表实现的集合。不接受 {@code null} 值。
 *
 * @param <C> 元素类型，要求有办法进行比较
 */
public class SkipListSet<C> implements Set<C> {
    /**
     * 塔高
     */
    private final int towerHeight;

    /**
     * 比较器。如果元素不是 {@link Comparable<C>}，就用它来比较。
     */
    private final Comparator<C> comparator;

    /**
     * 第一个塔。占位。
     */
    private final Tower firstTower;

    /**
     * 集合中元素数量。
     */
    private int size = 0;

    /**
     * 默认构造函数。
     */
    public SkipListSet() {
        this(null, DEFAULT_TOWER_HEIGHT);
    }

    /**
     * 创建时指定塔高。
     */
    public SkipListSet(int towerHeight) {
        this(null, towerHeight);
    }

    /**
     * 创建时指定比较器。
     */
    public SkipListSet(Comparator<C> comparator) {
        this(comparator, DEFAULT_TOWER_HEIGHT);
    }

    /**
     * 创建时指定比较器和塔高。
     */
    public SkipListSet(Comparator<C> comparator, int towerHeight) {
        this.comparator = comparator;
        this.towerHeight = Utils.fit(towerHeight, DEFAULT_TOWER_HEIGHT, MAX_TOWER_HEIGHT);
        this.firstTower = new Tower(ABSOLUTE_MIN, this.towerHeight);
        init();
    }

    private void init() {
        Tower lastTower = new Tower(ABSOLUTE_MAX, this.towerHeight);
        Arrays.fill(firstTower.next, lastTower);
        this.size = 0;
    }

    @Override
    public int size() {
        return this.size;
    }

    public void add(C element) {
        checkElement(element);

        Tower curr = this.firstTower;
        int level = this.towerHeight - 1;

        Tower[] temp = new Tower[this.towerHeight];

        while (level > -1) {
            final Tower nextTower = curr.next[level];
            final int result = compare(element, nextTower.value());
            if (result < 0) {
                temp[level] = curr;
                level--;
            } else if (result > 0) {
                curr = nextTower;
            } else {
                // 重复
                return;
            }
        }

        Tower newTower = new Tower(element, this.randomTowerHeight());
        for (int i = 0; i < newTower.next.length; i++) {
            newTower.next[i] = temp[i].next[i];
            temp[i].next[i] = newTower;
        }
        this.size++;
    }

    @Override
    public void clear() {
        init();
    }

    @Override
    public boolean has(C element) {
        checkElement(element);

        Tower curr = this.firstTower;
        int level = this.towerHeight - 1;

        while (level > -1) {
            final Tower nextTower = curr.next[level];
            final int result = compare(element, nextTower.value());
            if (result < 0) {
                level--;
            } else if (result > 0) {
                curr = nextTower;
            } else {
                return true;
            }
        }
        return false;
    }

    private int randomTowerHeight() {
        for (int i : Range.create(1, this.towerHeight - 1)) {
            if (Utils.randomInt(2) == 0) {
                return i;
            }
        }
        return this.towerHeight;
    }

    private void checkElement(C e) {
        Objects.requireNonNull(e);
        if (comparator == null && !(e instanceof Comparable)) {
            throw new IllegalArgumentException("Element should be Comparable");
        }
    }

    @SuppressWarnings("unchecked")
    private int naturallyCompare(C e1, C e2) {
        if (comparator != null) return comparator.compare(e1, e2);
        return ((Comparable<C>) e1).compareTo(e2);
    }

    private int compare(C e1, C e2) {
        if (e1 == ABSOLUTE_MAX) return 1;
        if (e1 == ABSOLUTE_MIN) return -1;

        if (e2 == ABSOLUTE_MAX) return -1;
        if (e2 == ABSOLUTE_MIN) return 1;

        return naturallyCompare(e1, e2);
    }

    @SuppressWarnings("unchecked")
    private final C ABSOLUTE_MIN = (C) new Object();

    @SuppressWarnings("unchecked")
    private final C ABSOLUTE_MAX = (C) new Object();


    private static final int DEFAULT_TOWER_HEIGHT = 8;
    private static final int MAX_TOWER_HEIGHT = 32;

    @Override
    public boolean remove(C element) {
        checkElement(element);

        Tower curr = this.firstTower;
        int level = this.towerHeight - 1;

        Tower[] temp = new Tower[this.towerHeight];
        Tower foundTower = null;
        while (level > -1) {
            final Tower nextTower = curr.next[level];
            final int result = compare(element, nextTower.value());
            if (result < 0) {
                temp[level] = curr;
                level--;
            } else if (result > 0) {
                curr = nextTower;
            } else {
                foundTower = nextTower;
                // 这里找到之后要继续往下找，直到找到最后一层
                temp[level] = curr;
                level--;
            }
        }

        if (foundTower == null) return false;

        for (int i : Range.create(foundTower.next.length - 1)) {
            temp[i].next[i] = foundTower.next[i];
        }

        this.size--;
        return true;
    }

    @Override
    public Iterator<C> iterator() {
        return new Itr();
    }

    private class Itr implements Iterator<C> {
        Tower curr = firstTower;

        @Override
        public boolean hasNext() {
            return curr.next[0].value != ABSOLUTE_MAX;
        }

        @Override
        public C next() {
            C result = curr.next[0].value();
            curr = curr.next[0];
            return result;
        }
    }

    private static class Tower {
        private final Object value;
        Tower[] next;

        @SuppressWarnings("unchecked")
        <E> E value() {
            return (E) value;
        }

        Tower(Object value, int towerHeight) {
            this.value = value;
            this.next = new Tower[towerHeight];
        }
    }
}
